from dd03_decoder import *


# 输出部分
class Generator(nn.Module):
    def __init__(self, d_model, vocab_size):
        super(Generator, self).__init__()
        self.d_model = d_model
        self.vocab_size = vocab_size

        # 定义线性输出层
        self.linear = nn.Linear(self.d_model, self.vocab_size)

    def forward(self, x):
        return F.log_softmax(self.linear(x), dim=-1)


def dm01_generator():
    """
    测试输出部分
    """
    # 准备解码器的输出结果
    decoder_result = dm02_decoder()

    # 实例化输出部分对象、调用输出层方法
    my_generator = Generator(d_model=512, vocab_size=1000)

    output = my_generator(decoder_result)

    print(f'输出层的输出结果:\n{output}')
    print(f'输出层的输出结果.shape:\n{output.shape}')


if __name__ == '__main__':
    dm01_generator()
